#include "hdr.h"

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif 

#define ERRIDX    -1
#define MINDIFF   0.00001
#define CHKFREQ   100


/* example on how to type cast pointer to a function 
typedef double (* deflocalf)(node_t *, int, double *);
typedef double (* defpairf)(node_t *, node_t *, int, int, double *);
*/

node_t *initall(int totalnodes, double **pval, int *pys, int *nbrnum, int **nbridx) 
{
   int      i, j;
   node_t   *p2allnodes, *p2nbr, *ptmp;

   printf("Start init all nodes...");

   p2allnodes = (node_t *) malloc(sizeof(node_t)*totalnodes);
   memset((char *)p2allnodes, 0, sizeof(node_t)*totalnodes);

   for(i=0;i<totalnodes;i++) {
      for(j=0;j<nbrnum[i];j++) {
         p2nbr = &p2allnodes[nbridx[i][j]];
         setneighbor(&p2allnodes[i], p2nbr);
      }
   }

   for(i=0;i<totalnodes;i++) {
      ptmp = &p2allnodes[i];
      if(nbrnum[i] != ptmp->nbrs) {
         printf("ERROR: neighbor data does not match!\n");
      }
      initnode(ptmp,i,pval[i],pys);
   }

   printf("Finish init all nodes\n");

   return p2allnodes;
}

/* find message value from box#boxid saved by psrc for y state
 */
double findmsg(node_t *pnode, node_t *psrc, int y, int boxid)
{
   int         i, status=0;
   double      msgval = 1;
   message_t   **mbox;
   message_t   *pmsg;

   if(boxid == BOXID1) {
      mbox = pnode->mbox1;
   } else {
      mbox = pnode->mbox2;
   }

   for(i=0;i<pnode->nbrs;i++) {
      pmsg = &mbox[y][i];
      if(pmsg->source == psrc) {
         msgval = pmsg->val;
         status = 1;
         break;
      }
   }

   if(!status) {
      printf("ERROR: cannot find matching message source!\n");
   }

   return msgval;
}

/* product of message transferred to pnode from all neighbors
 * except pnot with state of y from specific boxid
 */
double prodnbrmsgs(node_t *pnode, node_t *pnot, int y, int boxid)
{
   double      msgval = 1.0;
   node_t      *pcurnode;
   nodelist_t  *plist;

   plist = pnode->listhead;
   while(plist) {
      pcurnode = plist->pnode;
      if(pcurnode != pnot) {
         msgval = msgval * pow(findmsg(pnode, pcurnode, y, boxid),NE);
      }
      plist = plist->next;
   }

   return msgval;
}

/* find the normlization factor K for Mji(x_i)
 */
double normlizerK(node_t *pnode, node_t *psrc, int boxid)
{
   int      i;
   double   k, total = 0;

   for(i=0;i<NUMOFSTATE;i++) {
      total += findmsg(pnode, psrc, i, boxid);
   }
   
   k = 1 / total;
   return k;
}

/* send message from one node psrc to another node ptarget with state yi
 */
void sendmsg2(node_t *ptarget, node_t *psrc, int boxid, int yi, double msgval,
      int idx)
{
   message_t **mbox, *pmsg;
   int i;

   if(boxid == BOXID1) {
      mbox = ptarget->mbox1;
   } else {
      mbox = ptarget->mbox2;
   }

  for(i=0;i<ptarget->nbrs;i++) {
      pmsg = &mbox[yi][i];
      if(pmsg->source == psrc) {
         pmsg->val = msgval;
         pmsg->idx = idx;
         break;
      }
  }
}

/* transfer message from pnode to a neighbor for a specific state i
 */
void mp2node4i(node_t *psrc, node_t *ptarget, int yi, double *para1, 
      double *para2, int boxid)
{
   int         i;
   double      prodmsg=0.0, invmsg=0.0;
   double      e1=0.0, e2=0.0, ep=0.0, etotal=0.0;

   for(i=0;i<NUMOFSTATE;i++) {
      e1 = psrc->clocal(psrc, i, para1);
      e2 = psrc->cpair(psrc, ptarget, i, yi, para2);
      ep = exp(-1/NE*e2-e1);
      prodmsg = prodnbrmsgs(psrc, ptarget, i, boxid);
      invmsg = findmsg(psrc,ptarget,i,boxid);
      etotal += ep*prodmsg/(pow(invmsg,1-NE)+EPSILON);
   }
   etotal = etotal*normlizerK(ptarget, psrc, boxid);
   
   if(boxid == BOXID1) {
      //sendmsg2(ptarget, psrc, BOXID1, yi, etotal, -1);
      sendmsg2(ptarget, psrc, BOXID2, yi, etotal, -1);
   } else {
      //sendmsg2(ptarget, psrc, BOXID2, yi, etotal, -1);
      sendmsg2(ptarget, psrc, BOXID1, yi, etotal, -1);
   }
}

/* transfer message from pnode to a neighbor for a specific state i in 
 * max-product mode
 */
void mpmax2node4i(node_t *psrc, node_t *ptarget, int yi, double *para1, 
      double *para2, int boxid)
{
   int         i, idx;
   double      prodmsg=0.0, invmsg=0.0;
   double      e1=0.0, e2=0.0, ep=0.0, etotal=0.0, emax=0.0;

   for(i=0;i<NUMOFSTATE;i++) {
      e1 = psrc->clocal(psrc, i, para1);
      e2 = psrc->cpair(psrc, ptarget, i, yi, para2);
      ep = exp(-1/NE*e2-e1);
      prodmsg = prodnbrmsgs(psrc, ptarget, i, boxid);
      invmsg = findmsg(psrc,ptarget,i,boxid);
      etotal = ep*prodmsg/(pow(invmsg,1-NE)+EPSILON);
      if(etotal > emax) {
         emax = etotal;
         idx = i;
      }
   }
   emax = emax*normlizerK(ptarget, psrc, boxid);
   
   if(boxid == BOXID1) {
      //sendmsg2(ptarget, psrc, BOXID1, yi, emax, idx);
      sendmsg2(ptarget, psrc, BOXID2, yi, emax, idx);
   } else {
      //sendmsg2(ptarget, psrc, BOXID2, yi, emax, idx);
      sendmsg2(ptarget, psrc, BOXID1, yi, emax, idx);
   }
}


/* transfer message from every node to all of its neighbors for 
 * all states of y
 */
void runbpfd(node_t *phead, int totalnodes, double *para1, double *para2,
      int boxid)
{
   int         i, j;
   nodelist_t  *plist;

   for(i=0;i<totalnodes;i++) {
      plist = phead[i].listhead;
      while(plist) {
         for(j=0;j<NUMOFSTATE;j++) {
            mp2node4i(&phead[i], plist->pnode, j, para1, para2, boxid);
         }
         plist = plist->next;
      }
   }
}

void runbpbd(node_t *phead, int totalnodes, double *para1, double *para2,
      int boxid)
{
   int         i, j;
   nodelist_t  *plist;

   for(i=totalnodes-1;i>=0;i--) {
      plist = phead[i].listhead;
      while(plist) {
         for(j=0;j<NUMOFSTATE;j++) {
            mp2node4i(&phead[i], plist->pnode, j, para1, para2, boxid);
         }
         plist = plist->next;
      }
   }
}

/* transfer message from every node to all of its neighbors for 
 * all states of y in max product mode
 */
void runmaxbpfd(node_t *phead, int totalnodes, double *para1, double *para2,
      int boxid)
{
   int         i, j;
   nodelist_t  *plist;

   for(i=0;i<totalnodes;i++) {
      printf("Forward max-msg for node#%d\n", i); 
      plist = phead[i].listhead;
      while(plist) {
         for(j=0;j<NUMOFSTATE;j++) {
            mpmax2node4i(&phead[i], plist->pnode, j, para1, para2, boxid);
         }
         plist = plist->next;
      }
   }
}

void runmaxbpbd(node_t *phead, int totalnodes, double *para1, double *para2,
      int boxid)
{
   int         i, j;
   nodelist_t  *plist;

   for(i=totalnodes-1;i>=0;i--) {
      printf("Backward max-msg for node#%d\n", i);
      plist = phead[i].listhead;
      while(plist) {
         for(j=0;j<NUMOFSTATE;j++) {
            mpmax2node4i(&phead[i], plist->pnode, j, para1, para2, boxid);
         }
         plist = plist->next;
      }
   }
}

/* compute belief value bi for a node
 */
void computebi(node_t *pnode, double *para1, double *pk, int boxid)
{
   int         i, j;
   double      total=0.0, ei=0.0, prodmsg=1.0;
   message_t   **mbox;

   if(boxid == BOXID1) {
      mbox = pnode->mbox1;
   } else {
      mbox = pnode->mbox2;
   }
   
   for(i=0;i<NUMOFSTATE;i++) {
      prodmsg = 1.0;
      ei = exp(-pnode->clocal(pnode, i, para1));
      for(j=0;j<pnode->nbrs;j++) {
         prodmsg = prodmsg*pow(mbox[i][j].val, NE);
      }
      pnode->bi[i] = ei*prodmsg;
      total += pnode->bi[i];
   }
   for(i=0;i<NUMOFSTATE;i++) {
      *pk = 1 / total;
      pnode->bi[i] = pnode->bi[i] / total;
   }
}

/* compute pair belief bij for node i
 */

void computebij(node_t *pnode, double *para1, double *para2, double pk, 
      int boxid)
{
   int         i, j, t, idx=0;
   double      eij=0.0, ei=0.0, ej=0.0, prodmsg1=1.0, prodmsg2=1.0;
   double      invmsg1 = 1.0, invmsg2 = 1.0;
   nodelist_t  *plist;
   node_t      *pnbr;
   message_t   **mboxi, **mboxj;

   for(i=0;i<NUMOFSTATE;i++) {
      for(j=0;j<NUMOFSTATE;j++) {
         plist = pnode->listhead;
         idx = 0;
         while(plist) {
            pnbr = plist->pnode;
            if(boxid == BOXID1) {
               mboxi = pnode->mbox1;
               mboxj = pnbr->mbox1;
            } else {
               mboxi = pnode->mbox2;
               mboxj = pnbr->mbox2;
            }

            ei = pnode->clocal(pnode, i, para1);
            ej = pnbr->clocal(pnbr, j, para1);
            eij = pnode->cpair(pnode, pnbr, i, j, para2);
            for(t=0;t<pnode->nbrs;t++) {
               if(mboxi[i][t].source != pnbr) {
                  prodmsg1 *= pow(mboxi[i][t].val, NE);
               } else {
                  invmsg1 = pow(mboxi[i][t].val, 1-NE);
               }
            }
            for(t=0;t<pnbr->nbrs;t++) {
               if(mboxj[j][t].source != pnode) {
                  prodmsg2 *= pow(mboxj[j][t].val, NE);
               } else {
                  invmsg2 = pow(mboxj[j][t].val, 1-NE);
               }
            }
   
            pnode->bij[j+NUMOFSTATE*i][idx].val = pk*exp(-1/NE*eij-ei-ej) *
               prodmsg1/(invmsg1+EPSILON)*prodmsg2/(invmsg2+EPSILON);
            pnode->bij[j+NUMOFSTATE*i][idx].pnbr = pnbr;

            idx++;
            plist = plist->next;
         }
      }
   }
}

/* forward back tracing for the max-product result
 */
void fdbt4bp(node_t *phead, int totalnodes, int boxid)
{
   int         i, j, status, maxidx=0;
   double      maxbi=0.0;
   node_t      *pnode, *pnbr;
   nodelist_t  *plist;
   message_t   **mbox;

   printf("Start forward backtracing for blief value\n");

   pnode = &phead[0];
   for(i=0;i<NUMOFSTATE;i++) {
      if(pnode->bi[i] > maxbi) {
         maxbi = pnode->bi[i];
         maxidx = i;
      }
   }
   pnode->ymax = maxidx;

   for(i=0;i<totalnodes;i++) {
      pnode = &phead[i];
      if(pnode->ymax < 0) {
         printf("ERROR: the current node is not assigned in backtracing!\n");
      }  
      maxidx = pnode->ymax;

      if(boxid == BOXID1) {
         mbox = pnode->mbox1;
      } else {
         mbox = pnode->mbox2;
      }
  
      printf("Search forward max index for node#%d\n",i);
      plist = pnode->listhead;
      while(plist) {
         pnbr = plist->pnode;
         status = 0;
         for(j=0;j<pnode->nbrs;j++) {
            if(mbox[maxidx][j].source == pnbr) {
               status = 1;
               if(pnbr->ymax < 0) {
                  pnbr->ymax = mbox[maxidx][j].idx;
               }
               break;
            }
         }
         if(! status) {
            printf("ERROR: backtracing cannot find idx for %p\n", pnbr);
         }
         plist = plist->next;
      }
   }

   printf("Finish forward backtracing for belief value\n");
}

/* backward back tracing for the max-product result
 */
void bdbt4bp(node_t *phead, int totalnodes, int boxid)
{
   int         i, j, status, maxidx=0;
   double      maxbi=0.0;
   node_t      *pnode, *pnbr;
   nodelist_t  *plist;
   message_t   **mbox;

   printf("Start backward backtracing for blief value\n");
   pnode = &phead[totalnodes-1];
   for(i=0;i<NUMOFSTATE;i++) {
      if(pnode->bi[i] > maxbi) {
         maxbi = pnode->bi[i];
         maxidx = i;
      }
   }
   pnode->ymax = maxidx;

   for(i=totalnodes-1;i>=0;i--) {
      pnode = &phead[i];
      if(pnode->ymax < 0) {
         printf("ERROR: the current node is not assigned in backtracing!\n");
      }  
      maxidx = pnode->ymax;

      if(boxid == BOXID1) {
         mbox = pnode->mbox1;
      } else {
         mbox = pnode->mbox2;
      }
  
      printf("Search backward max index for node#%d\n",i);
      plist = pnode->listhead;
      while(plist) {
         pnbr = plist->pnode;
         status = 0;
         for(j=0;j<pnode->nbrs;j++) {
            if(mbox[maxidx][j].source == pnbr) {
               status = 1;
               if(pnbr->ymax < 0) {
                  pnbr->ymax = mbox[maxidx][j].idx;
               }
               break;
            }
         }
         if(! status) {
            printf("ERROR: backtracing cannot find idx for %p\n", pnbr);
         }
         plist = plist->next;
      }
   }

   printf("Finish backward backtracing for blief value\n");
}

nodelist_t *findlastinlist(nodelist_t *listhead)
{
   nodelist_t *prelist, *lastlist;
   prelist = listhead;
   lastlist = listhead;

   while(lastlist->next) {
      prelist = lastlist;
      lastlist = lastlist->next;
   }

   prelist->next = NULL;
   return lastlist;
}



double getlocal(node_t *pnode, int yi, double *para)
{
   nodelist_t *plist;
   node_t *pnbr;
   double val = 0.0;

   val = pnode->val[0];
   plist = pnode->listhead;
   while(plist) {
      pnbr = plist->pnode;
      /*val += fabs(pnode->val[0]-pnbr->val[0]);*/
      val += pnbr->val[0];
      plist = plist->next;
   }
   val = fabs(val/(pnode->nbrs+1) - yi);
   
   return val;
}

double getpair(node_t *pnode, node_t *pnbr, int ys, int yt, double *para)
{
   if(ys == yt)
      return 0.0;
   else
      return 1.0;
}

int chkconverge(node_t *phead, int *pys, int totalnodes)
{
   int i;
   double diffs = 0;

   for(i=0;i<totalnodes;i++) {
      diffs += abs(phead[i].ymax - pys[i]);
   }
   if(diffs < MINDIFF)
      return 1;
   else
      return 0;
}

void freeallnodes(node_t *phead, int totalnodes)
{
   int i, j;
   node_t *pnode;
   nodelist_t *plisthead, *p2del;

   printf("Start freeing mem for all nodes...");
   for(i=0;i<totalnodes;i++) {
      pnode = &phead[i];

      /* free nodelist from each node
       */
      plisthead = pnode->listhead;
      p2del = findlastinlist(plisthead);
      while(plisthead != p2del) {
         free(p2del);
         p2del = findlastinlist(plisthead);
      }
      free(plisthead);

      /* free message box from each node
       */
      for(j=0;j<NUMOFSTATE;j++) {
         free(pnode->mbox1[j]);
         free(pnode->mbox2[j]);
      }

      /* free bij from each node
       */
      for(j=0;j<NUMOFSTATE*NUMOFSTATE;j++) {
         free(pnode->bij[j]);
      }
   }

   free(phead);
   printf("Finish freeing mem for all nodes\n");
}


int main(int argc, char *argv[])
{
   int totalnodes=4, maxnbrs=4, maxxdims=1, maxloop=2;
   int pys[] ={0, 1};
   int pnbrnum[] = {2,2,2,2};
   int **pnbridx, i, boxid, pout[4];
   double **pxval, para1[1], para2[1], pk;
   node_t *phead;

    pnbridx = (int **) malloc(sizeof(int *)*totalnodes);
   for(i=0;i<totalnodes;i++) {
      pnbridx[i] = (int *) malloc(sizeof(int)*maxnbrs);
   }
   pxval = (double **) malloc(sizeof(double *)*totalnodes);
   for(i=0;i<totalnodes;i++) {
      pxval[i] = (double *) malloc(sizeof(double)*maxxdims);
   }

   /*
   pnbridx[0][0]=3;
   pnbridx[0][1]=1;
   pnbridx[0][2]=-1;
   pnbridx[0][3]=-1;
   pnbridx[1][0]=4;
   pnbridx[1][1]=2;
   pnbridx[1][2]=0;
   pnbridx[1][3]=-1;
   pnbridx[2][0]=5;
   pnbridx[2][1]=1;
   pnbridx[2][2]=-1;
   pnbridx[2][3]=-1;
   pnbridx[3][0]=6;
   pnbridx[3][1]=4;
   pnbridx[3][2]=0;
   pnbridx[3][3]=-1;
   pnbridx[4][0]=7;
   pnbridx[4][1]=5;
   pnbridx[4][2]=3;
   pnbridx[4][3]=1;
   pnbridx[5][0]=8;
   pnbridx[5][1]=4;
   pnbridx[5][2]=2;
   pnbridx[5][3]=-1;
   pnbridx[6][0]=7;
   pnbridx[6][1]=3;
   pnbridx[6][2]=-1;
   pnbridx[6][3]=-1;
   pnbridx[7][0]=8;
   pnbridx[7][1]=6;
   pnbridx[7][2]=4;
   pnbridx[7][3]=-1;
   pnbridx[8][0]=7;
   pnbridx[8][1]=5;
   pnbridx[8][2]=-1;
   pnbridx[8][3]=-1;
   */
   pnbridx[0][0]=2;
   pnbridx[0][1]=1;
   pnbridx[0][2]=-1;
   pnbridx[0][3]=-1;
   pnbridx[1][0]=3;
   pnbridx[1][1]=0;
   pnbridx[1][2]=-1;
   pnbridx[1][3]=-1;
   pnbridx[2][0]=3;
   pnbridx[2][1]=0;
   pnbridx[2][2]=-1;
   pnbridx[2][3]=-1;
   pnbridx[3][0]=2;
   pnbridx[3][1]=1;
   pnbridx[3][2]=-1;
   pnbridx[3][3]=-1;
 
   pxval[0][0] = 0.0;
   pxval[1][0] = 1.0;
   pxval[2][0] = 1.0;
   pxval[3][0] = 0.0;

   phead = initall(totalnodes, pxval, pys, pnbrnum, pnbridx);
   for(i=0;i<totalnodes;i++) {
      phead[i].clocal = getlocal;
      phead[i].cpair = getpair;
   }
   boxid = 1;

   for(i=0;i<maxloop;i++) {
      boxid = 1 - boxid;
      if((i % 2) == 0){
         printf("loop#%d: run forward max-msg passing\n", i);
         runmaxbpfd(phead, totalnodes, para1, para2, boxid);
      } else {
         printf("loop#%d, run backward max-msg passing\n", i);
         runmaxbpbd(phead, totalnodes, para1, para2, boxid);
      }
   }
   
   boxid = 1 - boxid;
   for(i=0;i<totalnodes;i++) {
      computebi(&phead[i], para1, &pk, boxid);
   }
   //bdbt4bp(phead, totalnodes, boxid);
   fdbt4bp(phead, totalnodes, boxid);
   for(i=0;i<totalnodes;i++) {
      pout[i] = (int) phead[i].ymax;
   }


   for(i=0;i<totalnodes;i++) {
      free(pnbridx[i]);
      free(pxval[i]);
   }
   free(pnbridx);
   free(pxval);
   freeallnodes(phead, totalnodes);

   return 0;
}
